// NetClient.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include <winsock2.h>
#include <ws2tcpip.h>
#include "../rdtsc.h"

#pragma comment (lib, "Ws2_32.lib")
#pragma comment (lib, "Mswsock.lib")
#pragma comment (lib, "AdvApi32.lib")

const __int64 cpuFreq = 2; //GHz
#define ITER 50000
#define ITER2 100
#define BUFLEN 131072

#define PORT 28100
#define STR_PORT "28100"


char buf[BUFLEN];
char servIPAddr[20] = "127.0.0.1";

void parse_args(int argc, char* argv[])
{
	int i = 1;
	while( i < argc ) {
		if(strcmp( argv[i], "-c") == 0 )
		{
			++i;
			if( i >= argc)
				break;
			strcpy_s(servIPAddr, argv[i]);
		}
		i++;
	}
}

int main(int argc, char* argv[])
{
	WSADATA wsaData;
	int iResult;
	SOCKET clientSocket = INVALID_SOCKET;
	__int64 t1, t2, total;
	total = 0;
	int recvBytes = 0;
	int recvTotal = ITER2 * BUFLEN;

	iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (iResult != NO_ERROR) {
        wprintf(L"WSAStartup function failed with error: %d\n", iResult);
        return 1;
    }	    

	parse_args(argc, argv);
	
	sockaddr_in servAddr;
    servAddr.sin_family = AF_INET;
    servAddr.sin_addr.s_addr = inet_addr(servIPAddr);
    servAddr.sin_port = htons(PORT);

	//setup socket
	t1 = rdtsc();
	clientSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (clientSocket == INVALID_SOCKET) {
        wprintf(L"socket function failed with error: %ld\n", WSAGetLastError());
        WSACleanup();
        return 1;
    }

	//connect to serv
	iResult = connect(clientSocket, (SOCKADDR *) & servAddr, sizeof (servAddr));
    if (iResult == SOCKET_ERROR) {
        wprintf(L"connect function failed with error: %ld\n", WSAGetLastError());
        iResult = closesocket(clientSocket);
        if (iResult == SOCKET_ERROR)
            wprintf(L"closesocket function failed with error: %ld\n", WSAGetLastError());
        WSACleanup();
        return 1;
    }
	t2 = rdtsc();
	double setupTime = (t2-t1) *1000 /(double)(cpuFreq*(1024*1024*1024));
    //wprintf(L"Connected to server.\n");
	
	BOOL bOptVal = TRUE;
	int bOptLen = sizeof(BOOL);
	//disable Nagle's algorithm for rtt
	iResult = setsockopt(clientSocket, IPPROTO_TCP, TCP_NODELAY, (char *) &bOptVal, bOptLen);	
	// test for round-trip time
	for(int i=0; i < ITER; ++i)
	{
		t1 = rdtsc();
		iResult = send( clientSocket, buf, 1, 0 );
		//if (iResult == SOCKET_ERROR) {
		//	printf("send failed with error: %d\n", WSAGetLastError());
		//	closesocket(clientSocket);
		//	WSACleanup();
		//	return 1;
		//}
		iResult = recv( clientSocket, buf, BUFLEN, 0 );
		t2 = rdtsc();
		if (iResult == SOCKET_ERROR) {
			printf("recv failed with error: %d\n", WSAGetLastError());
			closesocket(clientSocket);
			WSACleanup();
			return 1;
		}
		total += t2 - t1;
	}
	//send( clientSocket, buf, 0, 0 );
	//round trip time
	double rtt = (total * 1000 / (double)ITER) / (double)(cpuFreq*(1024*1024*1024)); //ms

	// enable Nagle's algorithm for bandwidth test
	bOptVal = FALSE;
	iResult = setsockopt(clientSocket, IPPROTO_TCP, TCP_NODELAY, (char *) &bOptVal, bOptLen);	

	// bandwidth test
	t1 = rdtsc();
	// upstream
    //#pragma omp parallel for
	for(int i=0; i < ITER2; ++i)
	{
		send( clientSocket, buf, BUFLEN, 0 );
	}
	t2 = rdtsc();

	//send( clientSocket, buf, 0, 0 );	
	// cal bandwidth
	__int64 bytesWritten = (__int64)ITER2 * (__int64)BUFLEN;
	double bandwidth1 = bytesWritten * cpuFreq / (double)(t2 - t1); // GB/s
	
	//downstream
	t1 = rdtsc();
	//#pragma omp parallel for
	while (recvBytes < recvTotal)
	{
		recvBytes += recv( clientSocket, buf, BUFLEN, 0 );
	}
	t2 = rdtsc();
	double bandwidth2 = recvBytes * cpuFreq / (double)(t2 - t1); // GB/s

	//shutdown
	t1 = rdtsc();
	iResult = shutdown(clientSocket, SD_BOTH);
	if (iResult == SOCKET_ERROR) {
		printf("shutdown failed with error: %d\n", WSAGetLastError());
		closesocket(clientSocket);
		WSACleanup();
		return 1;
	}
	iResult = closesocket(clientSocket);
	t2 = rdtsc();
	double tearDownTime = (t2-t1) *1000 /(double)(cpuFreq*(1024*1024*1024));

    WSACleanup();

	// report results
	printf("round-trip time: %.5f ms, bandwidth(upload): %.3f GB/s, bandwidth(download): %.3f GB/s, setup: %.5f ms, teardown: %.5f ms\n", rtt, bandwidth1, bandwidth2, setupTime, tearDownTime );
	return 0;
}

